/******************************************************************/
/*                                                                */
/*  Winpooch : Windows Watchdog                                   */
/*  Copyright (C) 2004-2007  Benoit Blanchon                      */
/*                                                                */
/*  This program is free software; you can redistribute it        */
/*  and/or modify it under the terms of the GNU General Public    */
/*  License as published by the Free Software Foundation; either  */
/*  version 2 of the License, or (at your option) any later       */
/*  version.                                                      */
/*                                                                */
/*  This program is distributed in the hope that it will be       */
/*  useful, but WITHOUT ANY WARRANTY; without even the implied    */
/*  warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR       */
/*  PURPOSE.  See the GNU General Public License for more         */
/*  details.                                                      */
/*                                                                */
/*  You should have received a copy of the GNU General Public     */
/*  License along with this program; if not, write to the Free    */
/*  Software Foundation, Inc.,                                    */
/*  675 Mass Ave, Cambridge, MA 02139, USA.                       */
/*                                                                */
/******************************************************************/


/******************************************************************/
/* Build configuration                                            */
/******************************************************************/

#define	TRACE_LEVEL	2


/******************************************************************/
/* Includes                                                       */
/******************************************************************/

// module's interface
#include "SymbolFile.h"

// project's headers
#include "Assert.h"
#include "Malloc.h"
#include "Trace.h"

#ifndef __NTDDK__
#include <ddk/ntddk.h>
#include <tchar.h>
#endif

/******************************************************************/
/* Internal macros                                                */
/******************************************************************/

#ifndef PACKED
#define PACKED __attribute__((__packed__))
#endif


/******************************************************************/
/* Internal constants                                             */
/******************************************************************/

#define SIGNATURE (*(DWORD*)"WP32")

#define FUNCTION_NAME_SIZE	32

#define VERSION_NAME_SIZE	64

#define SIGNATURE_SIZE		34

#define MAX_FUNCTION_COUNT	32

#define INITIAL_BUFFER_SIZE	(1<<20)

#define SIZE_INCREMENT		(1<<20)


/******************************************************************/
/* Internal data type                                             */
/******************************************************************/

typedef struct PACKED
{
  DWORD		dwSignature ;
  UCHAR		nFunctionCount ;
  UCHAR		nVersionNameSize ;
  UCHAR		nSignatureSize ;
  UCHAR		nFunctionNameSize ;
  DWORD		nFunctionTableOffset ;
  DWORD		nSymbolsOffset ;
} SYMFILEHEADER ;

typedef struct
{
  BYTE*		pBuffer ;
  UINT		nBufferTotalSize ;
  UINT		nBufferUsedSize ;
  UINT		nFunctionCount ;
  LPSTR		aFunctionNames[MAX_FUNCTION_COUNT] ;
} SYMFILE ;


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

HSYMFILE SymFile_NewW (UINT nFunctionCount, LPCWSTR aFuncNames[]) 
{
  HSYMFILE hSymFile ;
  int i ;
  LPCSTR* aFuncNamesA = (LPCSTR*) MALLOC(nFunctionCount*sizeof(LPSTR)) ;

  for( i=0 ; i<nFunctionCount ; i++ )
    {
      ANSI_STRING as ;
      UNICODE_STRING us ;
      int nSize ;

      RtlInitUnicodeString (&us, aFuncNames[i]) ;
      nSize = RtlUnicodeStringToAnsiSize(&us) ;
      
      as.Buffer = (LPSTR) MALLOC(nSize) ;
      as.MaximumLength = nSize ;
      as.Length = 0 ;          
      
      RtlUnicodeStringToAnsiString (&as, &us, FALSE) ;
      
      aFuncNamesA[i] = as.Buffer ;
    }

  hSymFile = SymFile_NewA (nFunctionCount, aFuncNamesA) ;

  for( i=0 ; i<nFunctionCount ; i++ )
    FREE ((void*)aFuncNamesA[i]) ;

  FREE (aFuncNamesA) ;

  return hSymFile ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

HSYMFILE SymFile_NewA (UINT nFunctionCount, LPCSTR* aFuncNames) 
{
  SYMFILE* pSymFile = (SYMFILE*) MALLOC (sizeof(SYMFILE)) ;
  SYMFILEHEADER *pHeader ;
  BYTE* pFuncNameInBuffer ;
  int i ;
  
  pSymFile->pBuffer = MALLOC (INITIAL_BUFFER_SIZE) ;
  pSymFile->nBufferTotalSize = INITIAL_BUFFER_SIZE ;
  pSymFile->nBufferUsedSize = 0 ;
  pSymFile->nFunctionCount = 0 ;

  pHeader = (SYMFILEHEADER*) pSymFile->pBuffer ;
  pHeader->dwSignature = SIGNATURE ;
  pHeader->nFunctionCount = nFunctionCount ;
  pHeader->nSignatureSize = SIGNATURE_SIZE ;
  pHeader->nVersionNameSize = VERSION_NAME_SIZE ;
  pHeader->nFunctionNameSize = FUNCTION_NAME_SIZE ;
  pHeader->nFunctionTableOffset = sizeof(SYMFILEHEADER) ;
  pHeader->nSymbolsOffset = sizeof(SYMFILEHEADER) + nFunctionCount*FUNCTION_NAME_SIZE ;
  pSymFile->nBufferUsedSize += sizeof(SYMFILEHEADER) ;

  pFuncNameInBuffer = pSymFile->pBuffer + pSymFile->nBufferUsedSize ;

  for( i=0 ; i<nFunctionCount ; i++ )
    {
      strcpy ((LPSTR)pFuncNameInBuffer, aFuncNames[i]) ;
      
      pSymFile->aFunctionNames[i] = pFuncNameInBuffer ;
      pSymFile->nFunctionCount++ ;

      pFuncNameInBuffer += pHeader->nFunctionNameSize ;
      pSymFile->nBufferUsedSize += pHeader->nFunctionNameSize ;
    }

  return (HSYMFILE) pSymFile ;
}


/******************************************************************/
/* Internal function                                              */
/******************************************************************/

HSYMFILE _SymFile_LoadUs (UNICODE_STRING *pusFilename) 
{
  SYMFILE*  pSymFile = (SYMFILE*) MALLOC (sizeof(SYMFILE)) ;
  SYMFILEHEADER *pHeader ;
  NTSTATUS nStatus ;
  HANDLE hFile ;
  OBJECT_ATTRIBUTES oaFileName ;
  IO_STATUS_BLOCK iostatus ;
  LARGE_INTEGER liSize ;
  LARGE_INTEGER liOffset ;
  FILE_STANDARD_INFORMATION fsi ;
  BYTE* pFuncNameInBuffer ;
  int i ;
  
  ASSERT (pusFilename!=NULL) ;
  ASSERT (pusFilename->Buffer!=NULL) ;



  TRACE_INFO (TEXT("Filename = %ls\n"), pusFilename->Buffer) ;

  InitializeObjectAttributes (&oaFileName, pusFilename, OBJ_CASE_INSENSITIVE, NULL, NULL) ;
		     
  liSize.QuadPart = pSymFile->nBufferUsedSize ;

  nStatus = ZwOpenFile (&hFile, GENERIC_READ|SYNCHRONIZE, &oaFileName, &iostatus, 0, FILE_SYNCHRONOUS_IO_NONALERT) ;

  if( nStatus != STATUS_SUCCESS )
    {
      TRACE_ERROR (TEXT("ZwOpenFile failed (status = 0x%08X)\n"), nStatus) ;
      FREE(pSymFile) ;
      return FALSE ;
    }
  
  nStatus = ZwQueryInformationFile (hFile, &iostatus, &fsi, sizeof(fsi), FileStandardInformation) ;

  if( nStatus != STATUS_SUCCESS )
    {
      TRACE_ERROR (TEXT("ZwQueryInformationFile failed (status = 0x%08X)\n"), nStatus) ;
      ZwClose (hFile) ;
      FREE(pSymFile) ;
      return FALSE ;
    }

  pSymFile->nBufferTotalSize = fsi.EndOfFile.QuadPart ;
  pSymFile->pBuffer = MALLOC(pSymFile->nBufferTotalSize) ;
  pSymFile->nBufferUsedSize = pSymFile->nBufferTotalSize ;
  pSymFile->nFunctionCount = 0 ;

  TRACE_INFO (TEXT("Reading %d bytes\n"), pSymFile->nBufferUsedSize) ;

  liOffset.QuadPart = 0 ; 
  nStatus = ZwReadFile (hFile, NULL, NULL, NULL, &iostatus, pSymFile->pBuffer, pSymFile->nBufferUsedSize, &liOffset, NULL) ;

  if( nStatus != STATUS_SUCCESS )
    {
      TRACE_ERROR (TEXT("ZwReadFile failed (status = 0x%08X)\n"), nStatus) ;
      ZwClose (hFile) ;
      FREE(pSymFile->pBuffer) ;
      FREE(pSymFile) ;
      return FALSE ;
    }

  ZwClose (hFile) ;

  pHeader = (SYMFILEHEADER*) pSymFile->pBuffer ;

  pFuncNameInBuffer = pSymFile->pBuffer + pHeader->nFunctionTableOffset ;

  for( i=0 ; i<pHeader->nFunctionCount ; i++ )
    {     
      TRACE_INFO(TEXT("Function%d : %hs\n"), i, pFuncNameInBuffer) ;

      pSymFile->aFunctionNames[i] = pFuncNameInBuffer ;
      pSymFile->nFunctionCount++ ;

      pFuncNameInBuffer += pHeader->nFunctionNameSize;
    }

  return (HSYMFILE)pSymFile ;  
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

HSYMFILE SymFile_LoadA (LPCSTR szFilename) 
{
  HSYMFILE hSymFile ;
  ANSI_STRING asFilename ;
  UNICODE_STRING usFilename ;

  RtlInitAnsiString (&asFilename, szFilename) ;

  RtlAnsiStringToUnicodeString (&usFilename, &asFilename, TRUE) ;
    
  hSymFile = _SymFile_LoadUs (&usFilename) ;

  RtlFreeUnicodeString (&usFilename) ;

  return hSymFile ;
}



/******************************************************************/
/* Internal function                                              */
/******************************************************************/

BOOL _SymFile_SaveUs (HSYMFILE hSymFile, UNICODE_STRING *pusFilename) 
{
  SYMFILE*  pSymFile = (SYMFILE*) hSymFile ;
  NTSTATUS nStatus ;
  HANDLE hFile ;
  OBJECT_ATTRIBUTES oaFileName ;
  IO_STATUS_BLOCK iostatus ;
  LARGE_INTEGER liSize ;
  LARGE_INTEGER liOffset ;
  
  ASSERT (hSymFile!=NULL) ;
  ASSERT (pusFilename!=NULL) ;
  ASSERT (pusFilename->Buffer!=NULL) ;

  TRACE_INFO (TEXT("Filename = %ls\n"), pusFilename->Buffer) ;

  InitializeObjectAttributes (&oaFileName, pusFilename, OBJ_CASE_INSENSITIVE, NULL, NULL) ;
		     
  liOffset.QuadPart = 0 ; 
  liSize.QuadPart = pSymFile->nBufferUsedSize ;

  nStatus = ZwCreateFile (&hFile, GENERIC_WRITE|SYNCHRONIZE, &oaFileName, &iostatus, &liSize, FILE_ATTRIBUTE_NORMAL, 0, FILE_SUPERSEDE, FILE_NON_DIRECTORY_FILE|FILE_SYNCHRONOUS_IO_NONALERT, NULL, 0) ;

  if( nStatus != STATUS_SUCCESS )
    {
      TRACE_ERROR (TEXT("ZwCreateFile failed (status = 0x%08X)\n"), nStatus) ;
      return FALSE ;
    }
  
  TRACE_INFO (TEXT("Writing %d bytes\n"), pSymFile->nBufferUsedSize) ;

  nStatus = ZwWriteFile (hFile, NULL, NULL, NULL, &iostatus, pSymFile->pBuffer, pSymFile->nBufferUsedSize, &liOffset, NULL) ;

  if( nStatus != STATUS_SUCCESS )
    {
      TRACE_ERROR (TEXT("ZwWriteFile failed (status = 0x%08X)\n"), nStatus) ;
      ZwClose (hFile) ;
      return FALSE ;
    }

  ZwClose (hFile) ;

  return TRUE ;  
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

BOOL SymFile_SaveW (HSYMFILE hSymFile, LPCWSTR szFilename) 
{
  UNICODE_STRING usFilename ;

  usFilename.Buffer = (LPWSTR)szFilename ;
  usFilename.MaximumLength = wcslen (szFilename) ;
  usFilename.Length = wcslen (szFilename) ;
  
  return _SymFile_SaveUs (hSymFile, &usFilename) ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

BOOL SymFile_SaveA (HSYMFILE hSymFile, LPCSTR szFilename) 
{
  ANSI_STRING asFilename ;
  UNICODE_STRING usFilename ;
  BOOL bResult ;

  RtlInitAnsiString (&asFilename, szFilename) ;

  RtlAnsiStringToUnicodeString (&usFilename, &asFilename, TRUE) ;
    
  bResult = _SymFile_SaveUs (hSymFile, &usFilename) ;

  RtlFreeUnicodeString (&usFilename) ;

  return bResult ;
}


/******************************************************************/
/* Internal function                                              */
/******************************************************************/

UINT _SymFile_GetKernelSize (SYMFILE* pSymFile)
{
  SYMFILEHEADER * pHeader ;

  ASSERT(pSymFile!=NULL) ;

  pHeader = (SYMFILEHEADER*) pSymFile->pBuffer ;

  return pHeader->nVersionNameSize + pHeader->nSignatureSize + pHeader->nFunctionCount * sizeof(INT_PTR) ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

HSYMKERN SymFile_AddKernel (HSYMFILE hSymFile) 
{
  SYMFILE*  pSymFile = (SYMFILE*) hSymFile ;
  UINT nKernSize ;
  UINT nKernOffset ;
    
  ASSERT(hSymFile!=NULL) ;

  nKernSize = _SymFile_GetKernelSize (hSymFile) ;

  while( pSymFile->nBufferUsedSize+nKernSize > pSymFile->nBufferTotalSize )
    {
      UINT nNewSize = pSymFile->nBufferTotalSize + SIZE_INCREMENT ;
      void* pNewBuffer = MALLOC(nNewSize) ;

      memcpy (pNewBuffer, pSymFile->pBuffer, pSymFile->nBufferUsedSize) ;

      pSymFile->pBuffer = pNewBuffer ;
      pSymFile->nBufferTotalSize = nNewSize ;
    }
  
  nKernOffset = pSymFile->nBufferUsedSize ;

  pSymFile->nBufferUsedSize += nKernSize ;

  return (HSYMKERN) nKernOffset ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

HSYMKERN SymFile_GetKernelA (HSYMFILE hSymFile, LPCSTR szSignature)
{
  INT_PTR nKernOffset ;
  INT_PTR nKernSize ;
  INT_PTR nKernSignOffset ;
  SYMFILE *pSymFile = (SYMFILE*)hSymFile ;
  SYMFILEHEADER *pHeader ;

  ASSERT (hSymFile!=NULL) ;
  ASSERT (szSignature!=NULL) ;

  pHeader = (SYMFILEHEADER*)pSymFile->pBuffer ;

  nKernOffset = pHeader->nSymbolsOffset ;
  nKernSize = _SymFile_GetKernelSize (hSymFile) ;
  nKernSignOffset = pHeader->nVersionNameSize ;

  while( nKernOffset < pSymFile->nBufferUsedSize )
    {
      LPCSTR szCurSign = (LPCSTR)( pSymFile->pBuffer + nKernOffset + nKernSignOffset ) ;
      
      TRACE_INFO (TEXT("Signature = %hs\n"), szCurSign) ;
 
      if( !strcmp(szCurSign, szSignature ) )
	return (HSYMFILE) nKernOffset ;
      
      nKernOffset += nKernSize ;
    }

  TRACE_WARNING (TEXT("Signature %hs is not present in symbol file\n"), szSignature) ;

  return NULL ;
}


/******************************************************************/
/* Internal function                                              */
/******************************************************************/

INT _SymFile_GetProcIndex (SYMFILE* pSymFile, LPCSTR szFuncName)
{
  int i ;
  
  ASSERT (pSymFile!=NULL) ;
  
  for( i=0 ; i<pSymFile->nFunctionCount ; i++ )
    if( ! strcmp(pSymFile->aFunctionNames[i], szFuncName) )
      return i ;
	
  TRACE_ERROR (TEXT("Function %hs not present in symbol file\n"), szFuncName) ;

  return -1 ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

INT_PTR SymFile_GetProcAddressA (HSYMFILE hSymFile, HSYMKERN hSymKern, LPCSTR szFuncName) 
{
  SYMFILE*	pSymFile = (SYMFILE*) hSymFile ;
  INT		iFuncIndex, nVersionNameSize, nSignatureSize ;
  INT_PTR*	pFuncAddrs ;

  ASSERT (hSymFile!=NULL) ;
  ASSERT (szFuncName!=NULL) ;
  ASSERT (pSymFile->pBuffer!=NULL) ;

  nVersionNameSize = ((SYMFILEHEADER*)pSymFile->pBuffer)->nVersionNameSize ;
  nSignatureSize = ((SYMFILEHEADER*)pSymFile->pBuffer)->nSignatureSize ;
  iFuncIndex = _SymFile_GetProcIndex (pSymFile, szFuncName) ;

  if( iFuncIndex < 0 )
    {
      TRACE_ERROR (TEXT("Function %hs is not present in symbol file\n"), szFuncName) ;
      return 0 ;
    }

  pFuncAddrs = (INT_PTR*)( (BYTE*)pSymFile->pBuffer + (INT_PTR)hSymKern + nVersionNameSize + nSignatureSize ) ;

  TRACE_INFO (TEXT("%hs(%d) => 0x%08X\n"), szFuncName, iFuncIndex, pFuncAddrs[iFuncIndex]) ;

  return pFuncAddrs[iFuncIndex] ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

BOOL SymFile_SetProcAddressA (HSYMFILE hSymFile, HSYMKERN hSymKern, LPCSTR szFuncName, INT_PTR nProcAddress) 
{
  SYMFILE*	pSymFile = (SYMFILE*) hSymFile ;
  INT		iFuncIndex, nVersionNameSize, nSignatureSize ;
  INT_PTR*	pFuncAddrs ;

  ASSERT (hSymFile!=NULL) ;
  ASSERT (szFuncName!=NULL) ;
  ASSERT (pSymFile->pBuffer!=NULL) ;

  nVersionNameSize = ((SYMFILEHEADER*)pSymFile->pBuffer)->nVersionNameSize ;
  nSignatureSize = ((SYMFILEHEADER*)pSymFile->pBuffer)->nSignatureSize ;
  iFuncIndex = _SymFile_GetProcIndex (pSymFile, szFuncName) ;

  if( iFuncIndex < 0 )
    {
      TRACE_ERROR (TEXT("Function %s is not present in symbol file\n"), szFuncName) ;
      return FALSE ;
    }

  pFuncAddrs = (INT_PTR*)( pSymFile->pBuffer + (INT_PTR)hSymKern + nVersionNameSize + nSignatureSize ) ;

  pFuncAddrs[iFuncIndex] = nProcAddress ;

  return TRUE ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

VOID SymFile_CloseFile (HSYMFILE hSymFile) 
{
  SYMFILE*	pSymFile = (SYMFILE*) hSymFile ;
  
  if( pSymFile != NULL )
    {
      if( pSymFile->pBuffer != NULL )
	FREE (pSymFile->pBuffer) ;

      FREE (pSymFile) ;
    }
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

BOOL SymFile_SetKernelNameA (HSYMFILE hSymFile, HSYMKERN hSymKern, LPCSTR szName)
{
  SYMFILE * pSymFile = (SYMFILE*) hSymFile ;
  INT_PTR nKernOffset = (INT_PTR)hSymKern ;
  BYTE* szNameInBuffer ;

  ASSERT (hSymFile!=NULL) ;
  ASSERT (hSymKern!=NULL) ;

  szNameInBuffer = pSymFile->pBuffer + nKernOffset ;
  strcpy ((LPSTR)szNameInBuffer, szName) ;

  return TRUE ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

BOOL SymFile_SetKernelSignatureA (HSYMFILE hSymFile, HSYMKERN hSymKern, LPCSTR szSignature)
{
  SYMFILE * pSymFile = (SYMFILE*) hSymFile ;
  INT_PTR nKernOffset = (INT_PTR)hSymKern ;
  BYTE* szSignInBuffer ;
  INT_PTR nKernSignOffset ;

  ASSERT (hSymFile!=NULL) ;
  ASSERT (hSymKern!=NULL) ;

  nKernSignOffset = ((SYMFILEHEADER*)pSymFile->pBuffer)->nVersionNameSize ;

  szSignInBuffer = pSymFile->pBuffer + nKernOffset + nKernSignOffset ;
  strcpy ((LPSTR)szSignInBuffer, szSignature) ;

  return TRUE ;
}
